from pymongo import MongoClient


class MongodbSource:
    def __init__(self, host, port, database):
        self.host = host
        self.port = port
        self.__init_database(database)

    def __del__(self):
        self.__del_database()

    def __init_database(self, database):
        self.client = MongoClient(self.host, self.port)
        self.database = self.client[database]

    def __del_database(self):
        try:
            self.client.close()
            self.client = None
            self.database = None
        except:
            pass

    def get_all_data(self, collection_name):
        collection = self.database[collection_name]
        for item in collection.find():
            yield item

    def insert_data(self, collection_name, data):
        collection = self.database[collection_name]
        collection.insert(data)
